#include <chrono>
#include <iostream>
#include <memory>
#include <numeric>
#include <opencv2/opencv.hpp>
#include <string>

#include "paddle_inference_api.h"

// using 在这里相当于声明，如果没有using的话就要用全名
using std::string;
using std::vector;
using std::cout;
using std::endl;
using paddle_infer::Config;
using paddle_infer::CreatePredictor;
using paddle_infer::Predictor;

// 初始化模型
std::shared_ptr<Predictor> InitPredictor(const std::string& model_file, const std::string& params_file) {
    Config config;
    config.SetModel(model_file, params_file);
    config.EnableMKLDNN();
    config.EnableMemoryOptim();
    return CreatePredictor(config);
}

// 模型推理
void run(Predictor *predictor,
         const std::vector<float> &input,
         const std::vector<int> &input_shape,
         std::vector<float> *out_data) {
    int input_num = std::accumulate(input_shape.begin(), input_shape.end(), 1, std::multiplies<int>());

    auto input_names = predictor->GetInputNames();
    auto output_names = predictor->GetOutputNames();
    auto input_t = predictor->GetInputHandle(input_names[0]);
    input_t->Reshape(input_shape);
    input_t->CopyFromCpu(input.data());

    predictor->Run();
    auto output_t = predictor->GetOutputHandle(output_names[0]);
    std::vector<int> output_shape = output_t->shape();
    int out_num = std::accumulate(output_shape.begin(), output_shape.end(), 1, std::multiplies<int>());
    out_data->resize(out_num);
    output_t->CopyToCpu(out_data->data());
}

// 读取图片数据
std::vector<float> read_img(std::string img_path) {
    cv::Mat img = cv::imread(img_path);

    // BGR 转 RGB
    cv::cvtColor(img, img, cv::COLOR_BGR2RGB);

    // Resize 到 224x224
    cv::resize(img, img, cv::Size(224, 224));

    // ResNet 预处理：减去均值 [0.485, 0.456, 0.406]，除以标准差 [0.229, 0.224, 0.225]
    std::vector<float> mean = {0.485, 0.456, 0.406};
    std::vector<float> std = {0.229, 0.224, 0.225};

    img.convertTo(img, CV_32F, 1.0 / 255.0);

    for (int c = 0; c < 3; ++c) {
        img.col(c) = img.col(c) - mean[c];
        img.col(c) = img.col(c) / std[c];
    }

    // 创建一个 vector<float> 来存储 CHW 格式的数据
    std::vector<float> img_data(224 * 224 * 3);

    // 将 HWC 转换为 CHW
    for (int c = 0; c < 3; ++c) {
        for (int h = 0; h < 224; ++h) {
            for (int w = 0; w < 224; ++w) {
                img_data[c * 224 * 224 + h * 224 + w] = (img.at<cv::Vec3f>(h, w)[c] - mean[c]) / std[c];
            }
        }
    }

    return img_data;
}

// 找到概率最大的那个类别的序号
int get_max_idx(const std::vector<float> &out_data) {
    int max_idx = 0;
    float max_val = out_data[0];
    for (size_t i = 1; i < out_data.size(); ++i) {
        if (out_data[i] > max_val) {
            max_idx = i;
            max_val = out_data[i];
        }
    }
    return max_idx;
}

int main(int argc, char *argv[]) {

    // 初始化模型预测器, 如果出现错误，请检查模型路径是否正确，最简单的做法就是使用绝对路径
    auto predictor = InitPredictor("../data/resnet50/inference.pdmodel" , "../data/resnet50/inference.pdiparams");

    // 读取图片数据，可以自由更换，input_shape代表输入数据的形状
    std::vector<int> input_shape = {1, 3, 224, 224};
    std::vector<float> input_data = read_img("data/dog.jpg");

    // 模型推理
    std::vector<float> out_data;
    run(predictor.get(), input_data, input_shape, &out_data);

    // 模型输出结果后处理，模型输出的内容是一个数组，这个数组记录了每一类别的概率，我们需要找到概率最大的那个类别
    int max_idx = get_max_idx(out_data);
    float max_val = out_data[max_idx];

    // 输出结果, 输出类别的序号，具体类别请访问下面的网址然后查看
    cout << "\n^^^^^^^the before output is the model structure, it is normal, do not worry^^^^^^^"<< endl;
    cout << "\n\n=================the output of the model is:============================"<< endl;
    cout << "max_idx: " << max_idx << " max_val: " << max_val << endl;
    cout << "you can explore https://blog.csdn.net/winycg/article/details/101722445 to know about the label of the output."<< endl;
    system("pause");
    return 0;
}